/***************************************************
Channel Coding Course Work: conolutional codes
This program template has given the message generator, BPSK modulation, AWGN channel model and BPSK demodulation,
you should first determine the encoder structure, then define the message and codeword length, generate the state table, write the convolutional encoder and decoder.
 
If you have any question, please contact me via e-mail: wuchy28@mail2.sysu.edu.cn
***************************************************/

#define _CRT_SECURE_NO_WARNINGS
#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <math.h>

#define message_length 1000 //the length of message

#define codeword_length 2000 //the length of codeword

float code_rate = (float)message_length / (float)codeword_length;

// channel coefficient
#define pi 3.1415926
double N0, sgm;

int state_table[8][5]; //state table, the size should be defined yourself

int reg_num = 2;

int state_num = 4 * 2; //the number of the state of encoder structure

int message[message_length];

int codeword[codeword_length];

int re_codeword[codeword_length]; //the received codeword
int de_message[message_length];	  //the decoding message

double tx_symbol[codeword_length][2]; //the transmitted symbols
double rx_symbol[codeword_length][2]; //the received symbols

//Variable using in hard decoder
int hard_branch_table[message_length][8];
int hard_path_table[message_length + 1][4];
int hard_ID_table[message_length][4];

//Variable using in soft decoder
double soft_branch_table[message_length][8];
double soft_path_table[message_length+1][4];
int soft_ID_table[message_length][4];




//Variable using in BCJR decoder
double forward_trace[message_length][4] = {0}; //Each node have a Prob
double backward_trace[message_length][4] = {0};
double transmit_table[message_length][8] = {0};
double channel_observations[codeword_length][2] = {0};

void statetable();
void encoder();
void modulation();
void demodulation();
void channel();
void decoder();


void hardDecoder();
void TableToInfinite();
void softDecoder();



//BCJR decoder function
void BCJRdecoder();
void channelObservations();
void forwardTrace();
void backwardTrace();
void bitsProb();

int main()
{
	int i;
	float SNR, start, finish;
	long int bit_error, seq, seq_num;
	double BER;
	double progress;

	//generate state table
	statetable();

	//random seed
	srand((int)time(0));

	//input the SNR and frame number
	printf("\nEnter start SNR: ");
	scanf("%f", &start);
	printf("\nEnter finish SNR: ");
	scanf("%f", &finish);
	printf("\nPlease input the number of message: ");
	scanf("%d", &seq_num);

	for (SNR = start; SNR <= finish; SNR++)
	{
		//channel noise
		N0 = (1.0 / code_rate) / pow(10.0, (float)(SNR) / 10.0);
		sgm = sqrt(N0 / 2);

		bit_error = 0;

		for (seq = 1; seq <= seq_num; seq++)
		{
			//generate binary message randomly
			/****************
			Pay attention that message is appended by 0 whose number is equal to the state of encoder structure.
			****************/
			for (i = 0; i < message_length - state_num; i++)
			{
				message[i] = rand() % 2;
			}
			for (i = message_length - state_num; i < message_length; i++)
			{
				message[i] = 0;
			}

			//convolutional encoder
			encoder();

			//BPSK modulation
			modulation();

			//AWGN channel
			channel();

			//BPSK demodulation, it's needed in hard-decision Viterbi decoder


			//硬判决注释下面
			demodulation();
			//hardDecoder();


			//软判决注释下面
			softDecoder();

			//BCJR注释下面
			//BCJRdecoder();


			//calculate the number of bit error
			for (i = 0; i < message_length; i++)
			{
				if (message[i] != de_message[i])
					bit_error++;
			}

			progress = (double)(seq * 100) / (double)seq_num;

			//calculate the BER
			BER = (double)bit_error / (double)(message_length * seq);

			//print the intermediate result
			printf("Progress=%2.1f, SNR=%2.1f, Bit Errors=%2.1d, BER=%E\r", progress, SNR, bit_error, BER);
		}

		//calculate the BER
		BER = (double)bit_error / (double)(message_length * seq_num);

		//print the final result
		printf("Progress=%2.1f, SNR=%2.1f, Bit Errors=%2.1d, BER=%E\n", progress, SNR, bit_error, BER);
	}
	system("pause");
	return 0;
}




/**
 * @brief 得到输入2bit数字的第一位或第二位
 * 
 * @param mode 1 为第一位，其他为第二位
 * @param num  输入的2bit数字
 * @return int 
 */
int bit(int mode, int num)
{
	return mode == 1 ? num / 10 : num % 10;
}

/**
 * @brief generate state_table
 * 
 */
void statetable()
{
	int S0 = 0;
	int S1 = 0;
	int current_state[8] = {00, 00, 01, 01, 10, 10, 11, 11}; //直接给值，和ppt上图一样

	for (int j = 0; j < state_num * 2; j++)
	{
		state_table[j][0] = j % 2;			  //input
		state_table[j][1] = current_state[j]; //Current state
		S0 = bit(1, state_table[j][1]);		  //put Current state into registers
		S1 = bit(2, state_table[j][1]);
		state_table[j][2] = S0 + j % 2 * 10;							   //Next state
		state_table[j][3] = (S0 + S1 + j % 2) % 2 * 10 + (j % 2 + S1) % 2; //OUT
		state_table[j][4] = j + 1;										   //ID, begin with 0
	}
}



/**
 * @brief Initialize table value to infinite，the value is codeword_length.
 * 
 */
void TableToInfinite()
{
	for (int i = 0; i < message_length; i++)
	{
		for (int j = 0; j < 4; j++)
		{

			soft_branch_table[i][2 * j] = codeword_length;
			soft_branch_table[i][2 * j + 1] = codeword_length;
			soft_path_table[i][j] = codeword_length;

			hard_branch_table[i][2 * j] = codeword_length;
			hard_branch_table[i][2 * j + 1] = codeword_length;
			hard_path_table[i][j] = codeword_length;
		}
	}
}


void hardDecoder()
{
	

	int current_state = 0;
	//calculate Branch metrics
	int state4[4] = {00, 01, 10, 11};
	//calculate path metrics
	int temp = message_length;
	int cnt;
//path table长度是message_length + 1，ID_table的是长度message_length，所以要用if来初始化
	for (int i = 0; i < message_length + 1; i++)//message_length的循环
	{
		if (i == 0)
		{
			soft_path_table[0][0] = 0;//初始化第一行，其他三个值为无穷
			
		}
		else
		{

			for (int j = 0; j < 4; j++)//节点的循环，分别是00，01，10，11
			{
				cnt = 0;//每次遍历state table 总会有两个满足条件，在两个branch中间选择更小的并且更新path metrics

				for (int m = 0; m < state_num; m++)//遍历state table
				{
					hard_branch_table[i][state_table[m][4] - 1] = abs(bit(1, state_table[m][3]) - re_codeword[i * 2]) + abs(bit(2, state_table[m][3]) - re_codeword[i * 2 + 1]);

					if (state_table[m][2] == state4[j])//当前节点等于state table 的next state ，这样可以找到上两条线是怎么连接的
					{
						cnt++;
						if (cnt == 1)//第一条线
						{
							temp = hard_branch_table[i - 1][state_table[m][4] - 1];  //记录下现在的branch metrics 为了和下一次的比较谁更小
							soft_path_table[i][j] = soft_path_table[i - 1][m / 2] + temp;//更新path table
							hard_ID_table[i - 1][j] = state_table[m][4];//记录ID
						}
						else if (cnt == 2)//第二条线到
						{
							if ((soft_path_table[i - 1][m / 2] + hard_branch_table[i - 1][state_table[m][4] - 1]) < soft_path_table[i][j])
							//比较第二条线和第一条线的path metrics 如果现在更小就更新成现在的，否则就不用更新
							{
								soft_path_table[i][j] = soft_path_table[i - 1][m / 2] + hard_branch_table[i - 1][state_table[m][4] - 1];//更新path table
								hard_ID_table[i - 1][j] = state_table[m][4];//更新ID table
							}
						}
					}
				}
			}
		}
	}

	int now_state;
	for (int i = message_length; i >= 0; i--)//遍历ID table 用回溯
	{
		if (i == message_length)
		{
			now_state = 0;//总是从最后一列第一行开始
		}
		else
		{
			de_message[i] = (hard_ID_table[i][now_state] - 1) % 2;
			now_state = (hard_ID_table[i][now_state] - 1) / 2;//比较巧妙的地方在于ID table里面的值减去1除2取整之后就是下一次的第几行
		}
	}
}

/**
 * @brief 计算欧式距离
 * 
 * @param c fhajshfksadhfk
 * @param c11 
 * @param c12 
 * @return int 
 */
double euDistance(int c, double c11, double c12)
{
	
	return (c == 0) ? sqrt((c11 - 1) * (c11 - 1) + (c12) * (c12)) : sqrt((c11 + 1) * (c11 + 1) + (c12) * (c12));
}


/**
 * @brief 软判决唯一改变的地方就是数组改为double，branch metrics改为欧式距离
 * 
 */
void softDecoder()
{
	

	for (int i = 0; i < message_length; i++)
	{
		for (int m = 0; m < state_num; m++)
		{
			soft_branch_table[i][state_table[m][4] - 1] = euDistance(bit(1, state_table[m][3]), rx_symbol[i * 2][0], rx_symbol[i * 2][1])+euDistance(bit(2, state_table[m][3]), rx_symbol[i * 2 + 1][0], rx_symbol[i * 2 + 1][1]);
		}

	}
	

	int current_state = 0;
	//calculate Branch metrics
	int state4[4] = {00, 01, 10, 11};
	//calculate path metrics
	double temp = message_length;
	int cnt;

	for (int i = 0; i < message_length + 1; i++)
	{
		if (i == 0)
		{
			soft_path_table[0][0] = 0;
			soft_path_table[0][1] = codeword_length;
			soft_path_table[0][2] = codeword_length;
			soft_path_table[0][3] = codeword_length;
		}
		else
		{

			for (int j = 0; j < 4; j++)
			{
				cnt = 0;

				for (int m = 0; m < state_num; m++)
				{

					if (state_table[m][2] == state4[j])
					{
						cnt++;
						if (cnt == 1)
						{
							temp = soft_branch_table[i - 1][state_table[m][4] - 1];
							soft_path_table[i][j] = soft_path_table[i - 1][m / 2] + temp;

							soft_ID_table[i - 1][j] = state_table[m][4];
						}
						else if (cnt == 2)
						{
							if ((soft_path_table[i - 1][m / 2] + soft_branch_table[i - 1][state_table[m][4] - 1]) < soft_path_table[i][j])
							{
								
								soft_path_table[i][j] = soft_path_table[i - 1][m / 2] + soft_branch_table[i - 1][state_table[m][4] - 1];

								soft_ID_table[i - 1][j] = state_table[m][4];
							}
						}
					}
				}
			}
		}
	}

	int now_state;
	for (int i = message_length; i >= 0; i--)
	{
		if (i == message_length)
		{
			now_state = 0;
		}
		else
		{
			de_message[i] = ( soft_ID_table[i][now_state] - 1 ) % 2;

			now_state = (soft_ID_table[i][now_state] - 1) / 2;
		}
	}
}






void BCJRdecoder()
{
	
	channelObservations();
	forwardTrace();
	backwardTrace();
	bitsProb();
}

void channelObservations()
{

	double Np;	 // normalization coefficient
	double dist; //OU distance
	for (int n = 0; n < codeword_length; n++)
	{
		Np = 0;
		//calculate p_ch(c_t^1=0)
		dist = pow(rx_symbol[n][0] - 1, 2) + pow(rx_symbol[n][1], 2);
		channel_observations[n][0] = (1 / (pi * N0)) * exp(-dist / N0);
		//calculate p_ch(c_t^1=1)
		dist = pow(rx_symbol[n][0] + 1, 2) + pow(rx_symbol[n][1], 2);
		channel_observations[n][1] = (1 / (pi * N0)) * exp(-dist / N0);
		Np = channel_observations[n][0] + channel_observations[n][1];

		// normalization
		channel_observations[n][0] = channel_observations[n][0] / Np;
		channel_observations[n][1] = channel_observations[n][1] / Np;
	}

	for (int m = 0; m < message_length; m++) //obtain from Vertibi terllis
	{

		transmit_table[m][0] = channel_observations[2 * m][0] * channel_observations[2 * m + 1][0];
		transmit_table[m][1] = channel_observations[2 * m][1] * channel_observations[2 * m + 1][1];
		transmit_table[m][2] = transmit_table[m][1];
		transmit_table[m][3] = transmit_table[m][0];
		transmit_table[m][4] = channel_observations[2 * m][1] * channel_observations[2 * m + 1][0];
		transmit_table[m][5] = channel_observations[2 * m][0] * channel_observations[2 * m + 1][1];
		transmit_table[m][6] = transmit_table[m][5];
		transmit_table[m][7] = transmit_table[m][4];

		Np = transmit_table[m][0] + transmit_table[m][1] + transmit_table[m][2] + transmit_table[m][3] + transmit_table[m][4] + transmit_table[m][5] + transmit_table[m][6];

		for (int j = 0; j < 8; j++)
		{
			transmit_table[m][j] = transmit_table[m][j] / Np;
		}
	}
}
void forwardTrace()
{
	double Np = 0;
	for (int i = 0; i < message_length; i++)
	{

		if (i == 0)
		{
			forward_trace[i][0] = 1; //Initialize probability of state00=1,other is 0
			forward_trace[i][1] = 0;
			forward_trace[i][2] = 0;
			forward_trace[i][3] = 0;
		}
		else
		{
			forward_trace[i][0] = transmit_table[i - 1][0] * forward_trace[i - 1][0] + transmit_table[i - 1][2] * forward_trace[i - 1][1]; //00->00+11->00
			forward_trace[i][1] = transmit_table[i - 1][4] * forward_trace[i - 1][2] + transmit_table[i - 1][6] * forward_trace[i - 1][3]; //10->01+11->01
			forward_trace[i][2] = transmit_table[i - 1][1] * forward_trace[i - 1][0] + transmit_table[i - 1][3] * forward_trace[i - 1][1]; //00->10+01->10
			forward_trace[i][3] = transmit_table[i - 1][5] * forward_trace[i - 1][2] + transmit_table[i - 1][7] * forward_trace[i - 1][3]; //10->11+11->11
		}
		//normalize
		Np = forward_trace[i][0] + forward_trace[i][1] + forward_trace[i][2] + forward_trace[i][3];
		forward_trace[i][0] = forward_trace[i][0] / Np;
		forward_trace[i][1] = forward_trace[i][1] / Np;
		forward_trace[i][2] = forward_trace[i][2] / Np;
		forward_trace[i][3] = forward_trace[i][3] / Np;
	}
}
void backwardTrace()
{

	double Np = 0;
	for (int i = message_length - 1; i >= 0; i--)
	{

		if (i == message_length - 1)
		{
			backward_trace[i][0] = 1; //Initialize probability of state00=1,other is 0
			backward_trace[i][1] = 0;
			backward_trace[i][2] = 0;
			backward_trace[i][3] = 0;
		}
		else
		{
			backward_trace[i][0] = transmit_table[i + 1][0] * backward_trace[i + 1][0] + transmit_table[i + 1][1] * backward_trace[i + 1][2]; //00<-00+00<-10
			backward_trace[i][1] = transmit_table[i + 1][2] * backward_trace[i + 1][0] + transmit_table[i + 1][3] * backward_trace[i + 1][2]; //01<-00+01<-10
			backward_trace[i][2] = transmit_table[i + 1][4] * backward_trace[i + 1][1] + transmit_table[i + 1][5] * backward_trace[i + 1][3]; //10<-01+10<-11
			backward_trace[i][3] = transmit_table[i + 1][6] * backward_trace[i + 1][1] + transmit_table[i + 1][7] * backward_trace[i + 1][3]; //11<-11+11<-01
		}
		//normalize

		Np = backward_trace[i][0] + backward_trace[i][1] + backward_trace[i][2] + backward_trace[i][3];
		backward_trace[i][0] = backward_trace[i][0] / Np;
		backward_trace[i][1] = backward_trace[i][1] / Np;
		backward_trace[i][2] = backward_trace[i][2] / Np;
		backward_trace[i][3] = backward_trace[i][3] / Np;
	}
}

void bitsProb()
{
	double p0 = 0;
	double p1 = 0;
	for (int i = 0; i < message_length; i++)
	{ //p0=0 0 0 +1 2 0 +2 4 1+ 3 6 1 第一位数是起始状态，从trellis第一行往下数 中间是状态转移 
		p0 = forward_trace[i][0] * transmit_table[i][0] * backward_trace[i][0] + forward_trace[i][1] * transmit_table[i][2] * backward_trace[i][0] + forward_trace[i][2] * transmit_table[i][4] * backward_trace[i][1] + forward_trace[i][3] * transmit_table[i][6] * backward_trace[i][1];
		//p1=0 1 2+1 3 2+ 2 5 3+3 7 3
		p1 = forward_trace[i][0] * transmit_table[i][1] * backward_trace[i][2] + forward_trace[i][1] * transmit_table[i][3] * backward_trace[i][2] + forward_trace[i][2] * transmit_table[i][5] * backward_trace[i][3] + forward_trace[i][3] * transmit_table[i][7] * backward_trace[i][3];
		de_message[i] = p0 > p1 ? 0 : 1;
	}
}

/**
 * @brief Using state_table/generation matrix  obains codework (7,5)8 conv-code
 * 
 */
void encoder()
{
	//convolution encoder, the input is message[] and the output is codeword[]
	int current_state = 0;

	for (int i = 0; i < message_length; i++)
	{
		for (int j = 0; j < state_num; j++)
		{
			if ((message[i] == state_table[j][0]) & (current_state == state_table[j][1]))//if  input and current state are equal to state_table
			{
				current_state = state_table[j][2]; //update current_state to next_state,obtian OUT

				codeword[2 * i] = bit(1, state_table[j][3]);	 //c1
				codeword[2 * i + 1] = bit(2, state_table[j][3]); //c2
				break;
			}
		}
	}

	/*用生成矩阵的方式
	int message[message_length] = {1, 0, 1, 1, 0};
    int reg1[3] = {1, 1, 1};
    int reg2[3] = {1, 0, 1};
    int G_part1[message_length][codeword_length] = {0};
    int G_part2[message_length][codeword_length] = {0};
    int G[message_length][codeword_length] = {0};
    for (int i = 0; i < message_length; i++)
    {
        for (int j = 0; j < codeword_length; j++)
        {
            if (j - i < 3 & j - i >= 0)
            {
                G_part1[i][j] = reg1[(j - i) % 3];
                G_part2[i][j] = reg2[(j - i) % 3];
            }
            G[i][2 * j] = G_part1[i][j];
            G[i][2 * j + 1] = G_part2[i][j];
        }
    }
    for (int i = 0; i < codeword_length; i++)
    {
        for (int j = 0; j < message_length; j++)
        {
            codeword[i] = (codeword[i] + (G[j][i] * message[j]))%2;//矩阵相乘
        }
        //return 0;
    }
	*/
}

void modulation()
{
	//BPSK modulation
	int i;

	//0 is mapped to (1,0) and 1 is mapped tp (-1,0)
	for (i = 0; i < codeword_length; i++)
	{
		tx_symbol[i][0] = -1 * (2 * codeword[i] - 1);
		tx_symbol[i][1] = 0;
	}
}

void channel()
{
	//AWGN channel
	int i, j;
	double u, r, g;

	for (i = 0; i < codeword_length; i++)
	{
		for (j = 0; j < 2; j++)
		{
			u = (float)rand() / (float)RAND_MAX;
			if (u == 1.0)
				u = 0.999999;
			r = sgm * sqrt(2.0 * log(1.0 / (1.0 - u)));

			u = (float)rand() / (float)RAND_MAX;
			if (u == 1.0)
				u = 0.999999;
			g = (float)r * cos(2 * pi * u);

			rx_symbol[i][j] = tx_symbol[i][j] + g;
		}
	}
}
void demodulation()
{
	int i;
	double d1, d2;
	for (i = 0; i < codeword_length; i++)
	{
		d1 = (rx_symbol[i][0] - 1) * (rx_symbol[i][0] - 1) + rx_symbol[i][1] * rx_symbol[i][1];
		d2 = (rx_symbol[i][0] + 1) * (rx_symbol[i][0] + 1) + rx_symbol[i][1] * rx_symbol[i][1];
		if (d1 < d2)
			re_codeword[i] = 0;
		else
			re_codeword[i] = 1;
	}
}






